# _*_ coding: utf-8 _*_
# ---------------------------
# @创建作者：ZQQ
# @创建时间：2025/01/15
# @说明：任务状态管理模块
# ---------------------------

import threading
import time
import uuid
from datetime import datetime
from typing import Dict, Optional


class TaskStatus:
    """任务状态常量"""
    PENDING = "pending"
    DOWNLOADING = "downloading"
    BUILDING = "building"
    STARTING = "starting"
    SUCCESS = "success"
    FAILED = "failed"


class DeploymentTask:
    """部署任务类"""
    
    def __init__(self, task_id: str, model_name: str, port: int, service_folder: str):
        """
        初始化部署任务
        
        Args:
            task_id: 任务ID
            model_name: 模型名称
            port: 端口号
            service_folder: 服务文件夹名称
        """
        self.task_id = task_id
        self.model_name = model_name
        self.port = port
        self.service_folder = service_folder
        self.status = TaskStatus.PENDING
        self.message = "任务已创建，等待开始..."
        self.progress = 0
        self.created_at = datetime.now()
        self.updated_at = datetime.now()
        self.result = None
        self.error = None
    
    def update_status(self, status: str, message: str, progress: Optional[int] = None):
        """
        更新任务状态
        
        Args:
            status: 新状态
            message: 状态消息
            progress: 进度百分比（可选）
        """
        with task_lock:
            self.status = status
            self.message = message
            if progress is not None:
                self.progress = progress
            self.updated_at = datetime.now()
    
    def set_result(self, success: bool, result: Optional[Dict] = None, error: Optional[str] = None):
        """
        设置任务结果
        
        Args:
            success: 是否成功
            result: 成功结果（可选）
            error: 错误信息（可选）
        """
        with task_lock:
            self.status = TaskStatus.SUCCESS if success else TaskStatus.FAILED
            self.result = result
            self.error = error
            self.updated_at = datetime.now()
            if success:
                self.message = "部署成功完成"
                self.progress = 100
            else:
                self.message = f"部署失败: {error}"
    
    def to_dict(self) -> Dict:
        """
        转换为字典格式
        
        Returns:
            任务信息的字典表示
        """
        return {
            'taskId': self.task_id,
            'status': self.status,
            'message': self.message,
            'progress': self.progress,
            'createdAt': self.created_at.isoformat(),
            'updatedAt': self.updated_at.isoformat(),
            'modelName': self.model_name,
            'port': self.port,
            'serviceFolder': self.service_folder,
            'result': self.result,
            'error': self.error
        }


class TaskManager:
    """任务管理器"""
    
    def __init__(self):
        """初始化任务管理器"""
        self.tasks: Dict[str, DeploymentTask] = {}
        self.lock = threading.Lock()
    
    def create_task(self, model_name: str, port: int, service_folder: str) -> str:
        """
        创建新的部署任务
        
        Args:
            model_name: 模型名称
            port: 端口号
            service_folder: 服务文件夹名称
            
        Returns:
            任务ID
        """
        task_id = str(uuid.uuid4())
        task = DeploymentTask(task_id, model_name, port, service_folder)
        
        with self.lock:
            self.tasks[task_id] = task
        
        return task_id
    
    def get_task(self, task_id: str) -> Optional[DeploymentTask]:
        """
        获取任务
        
        Args:
            task_id: 任务ID
            
        Returns:
            任务对象，如果不存在则返回None
        """
        with self.lock:
            return self.tasks.get(task_id)
    
    def get_all_tasks(self) -> Dict[str, DeploymentTask]:
        """
        获取所有任务
        
        Returns:
            所有任务的字典
        """
        with self.lock:
            return self.tasks.copy()
    
    def list_tasks(self) -> list:
        """
        列出所有任务的基本信息
        
        Returns:
            任务信息列表
        """
        with self.lock:
            return [task.to_dict() for task in self.tasks.values()]
    
    def remove_task(self, task_id: str) -> bool:
        """
        删除任务
        
        Args:
            task_id: 任务ID
            
        Returns:
            是否删除成功
        """
        with self.lock:
            if task_id in self.tasks:
                del self.tasks[task_id]
                return True
            return False
    
    def cleanup_old_tasks(self, max_age_hours: int = 24):
        """
        清理旧任务
        
        Args:
            max_age_hours: 最大保留时间（小时）
        """
        current_time = datetime.now()
        max_age_seconds = max_age_hours * 3600
        
        with self.lock:
            tasks_to_remove = []
            for task_id, task in self.tasks.items():
                age_seconds = (current_time - task.created_at).total_seconds()
                if age_seconds > max_age_seconds:
                    tasks_to_remove.append(task_id)
            
            for task_id in tasks_to_remove:
                del self.tasks[task_id]
        
        if tasks_to_remove:
            print(f"🧹 清理了 {len(tasks_to_remove)} 个旧任务")


# 全局任务管理器实例
task_manager = TaskManager()

# 全局锁（为了向后兼容）
task_lock = task_manager.lock
